from core.cuda import cuda
from core.config import no_grad
from core.data.dataloader import DataLoader
from core.trainer import Trainer
from implement.datasource.minist import MNIST
from implement.models.mlp import MLP
from implement.optimizers.adam import Adam
from implement.optimizers.hooks.weight_decay import WeightDecay
from utils.common import accuracy
from utils.functions_collect import softmax_cross_entropy, relu

if '__file__' in globals():
    import os, sys

    sys.path.append(os.path.join(os.path.dirname(__file__), '..'))

max_epoch = 5
batch_size = 100
hidden_size = 1000

train_set = MNIST(train=True)
test_set = MNIST(train=False)

model = MLP((hidden_size, hidden_size, 10), activation=relu)
optimizer = Adam().setup(model)
optimizer.add_hook(WeightDecay(1e-4))  # Weight decay

train = Trainer(model, optimizer)
train.prepare_data(train_set, test_set)
train.load_weights('mnist_weights.npz')
train.fit(softmax_cross_entropy,max_epoch)
train.save_weights('mnist_weights.npz')
train.plot()
